import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import torchvision

# 设置随机种子，确保结果可复现
# torch.manual_seed(42)

# 设备配置，检查是否有可用的GPU，否则使用CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 超参数设置
batch_size = 64  # 每批次样本数
lr = 0.0002  # 学习率
z_dim = 100  # 随机噪声的维度
epochs = 50  # 训练轮数

# 数据加载和预处理，进行标准化处理
transform = transforms.Compose(
    [
        transforms.ToTensor(),  # 将图像转换为Tensor
        transforms.Normalize(mean=[0.5], std=[0.5]),  # 将像素值归一化到[-1, 1]范围
    ]
)

# 下载MNIST数据集，应用预处理
train_data = datasets.MNIST(
    root="./data", train=True, download=True, transform=transform
)

# 使用DataLoader加载数据集
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)


# 生成器（Generator）类定义
class Generator(nn.Module):
    def __init__(self, z_dim):
        super(Generator, self).__init__()
        # 定义生成器的全连接层网络结构
        self.fc = nn.Sequential(
            nn.Linear(z_dim, 256),  # 输入随机噪声，维度为z_dim，输出256维
            nn.LeakyReLU(0.2),  # 使用LeakyReLU激活函数
            nn.Linear(256, 512),  # 从256维转换到512维
            nn.LeakyReLU(0.2),  # 使用LeakyReLU激活函数
            nn.Linear(512, 1024),  # 从512维转换到1024维
            nn.LeakyReLU(0.2),  # 使用LeakyReLU激活函数
            nn.Linear(1024, 28 * 28),  # 最后输出28*28的图像（MNIST图像的大小）
            nn.Tanh(),  # 使用Tanh激活函数，将输出范围压缩到[-1, 1]
        )

    def forward(self, z):
        # 输入噪声z，输出生成的图像
        return self.fc(z).view(z.size(0), 1, 28, 28)  # 将输出展平为28x28的图像


# 判别器（Discriminator）类定义
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        # 定义判别器的全连接层网络结构
        self.fc = nn.Sequential(
            nn.Linear(28 * 28, 1024),  # 输入图像展平为一维向量
            nn.LeakyReLU(0.2),  # 使用LeakyReLU激活函数
            nn.Linear(1024, 512),  # 从1024维转换到512维
            nn.LeakyReLU(0.2),  # 使用LeakyReLU激活函数
            nn.Linear(512, 256),  # 从512维转换到256维
            nn.LeakyReLU(0.2),  # 使用LeakyReLU激活函数
            nn.Linear(256, 1),  # 输出一个标量，表示图像的真实性
            nn.Sigmoid(),  # 使用Sigmoid激活函数，将输出压缩到[0, 1]区间，表示真假概率
        )

    def forward(self, img):
        # 输入图像，返回判别结果（真假）
        return self.fc(img.view(img.size(0), -1))  # 展平图像并通过判别器


# 初始化生成器和判别器
generator = Generator(z_dim).to(device)
discriminator = Discriminator().to(device)

# 损失函数（二进制交叉熵）和优化器（Adam优化器）
criterion = nn.BCELoss()  # 使用二进制交叉熵损失
optimizer_g = optim.Adam(
    generator.parameters(), lr=lr, betas=(0.5, 0.999)
)  # 生成器的优化器
optimizer_d = optim.Adam(
    discriminator.parameters(), lr=lr, betas=(0.5, 0.999)
)  # 判别器的优化器

# 训练过程
for epoch in range(epochs):  # 遍历每个epoch
    for i, (imgs, _) in enumerate(train_loader):  # 加载每个batch
        imgs = imgs.to(device)  # 将图片数据传入设备（GPU/CPU）
        batch_size = imgs.size(0)  # 获取当前batch的大小

        # 真实标签和假标签
        real_labels = torch.ones(batch_size, 1).to(device)  # 真实图片标签是1
        fake_labels = torch.zeros(batch_size, 1).to(device)  # 假图片标签是0

        # 判别器训练
        optimizer_d.zero_grad()  # 清空判别器的梯度

        # 计算真实图片的损失
        real_imgs = imgs.view(batch_size, -1)  # 展平输入图片
        output_real = discriminator(real_imgs)  # 判别器判断真实图片
        d_loss_real = criterion(output_real, real_labels)  # 计算真实图片的损失

        # 生成假图片并计算损失
        z = torch.randn(batch_size, z_dim).to(device)  # 从标准正态分布生成随机噪声
        fake_imgs = generator(z)  # 通过生成器生成假图像
        fake_imgs = fake_imgs.view(batch_size, -1)  # 展平生成的假图像
        output_fake = discriminator(
            fake_imgs.detach()
        )  # 判别器判断假图像（不计算梯度）
        d_loss_fake = criterion(output_fake, fake_labels)  # 计算假图片的损失

        # 总损失并反向传播
        d_loss = d_loss_real + d_loss_fake  # 判别器的总损失
        d_loss.backward()  # 反向传播计算梯度
        optimizer_d.step()  # 更新判别器的权重

        # 生成器训练
        optimizer_g.zero_grad()  # 清空生成器的梯度

        output_fake = discriminator(fake_imgs)  # 判别器判断假图像
        g_loss = criterion(
            output_fake, real_labels
        )  # 生成器希望判别器认为假图像是真实的
        g_loss.backward()  # 反向传播计算梯度
        optimizer_g.step()  # 更新生成器的权重

        # 打印每100步的进度
        if i % 100 == 0:
            print(
                f"Epoch [{epoch}/{epochs}], Step [{i}/{len(train_loader)}], "
                f"D Loss: {d_loss.item():.4f}, G Loss: {g_loss.item():.4f}"
            )

    # 每个epoch结束后，生成一些图像并展示
    with torch.no_grad():  # 在生成图像时不计算梯度
        z = torch.randn(64, z_dim).to(device)  # 生成64个随机噪声样本
        fake_images = generator(z)  # 使用生成器生成假图像
        fake_images = fake_images.cpu().data  # 转移到CPU并取数据
        grid = torchvision.utils.make_grid(
            fake_images, nrow=8, normalize=True
        )  # 将生成的图像拼接成网格
        plt.imshow(grid.permute(1, 2, 0))  # 显示图像
        # plt.show()  # 展示图像
        plt.savefig(
            os.path.join(os.path.dirname(__file__), f"epoch_{epoch}_fake_img.png")
        )

# 训练结束后保存模型
torch.save(generator.state_dict(), "generator.pth")  # 保存生成器模型
torch.save(discriminator.state_dict(), "discriminator.pth")  # 保存判别器模型
